from scipy import stats
import numpy as np
import sys
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def test(N,times):
    S=np.zeros((N,4))
    for i in range(N):
        for j in range(times):
            t=np.random.rand()
            if t<1/6:
                S[i][0]+=1
            elif t>=1/6 and t<1/3:
                S[i][0]-=1
            elif t>=1/3 and t<1/2:
                S[i][1]+=1
            elif t>=1/2 and t<2/3:
                S[i][1]-=1
            elif t>=2/3 and t<5/6:
                S[i][2]+=1
            else:
                S[i][2]-=1
            if dis_3(S[i])==0:
                S[i][3]=1
    return S
def pro_test(S):
    n=len(S)
    k=0
    for i in range(n):
        if S[i][3]==1:
            k+=1
    return k/n
def dis_3(X):
    return np.sqrt(X[0]*X[0]+X[1]*X[1]+X[2]*X[2])           
def dis(S):
    n=len(S)
    A=np.zeros(n)
    for i in range(n):
        A[i]=dis_3(S[i])
    return A    
def print_test(S):
    A=dis(S)
    fig=plt.figure(1)
    plt.hist(A, bins = 30)
    plt.xlabel("distance")
    plt.ylabel("number of point")
    plt.show()
def draw_test(S):
    fig=plt.figure(2)
    ax = Axes3D(fig)
    x=S[...,0]
    y=S[...,1]
    z=S[...,2]
    ax.scatter3D(x,y,z, cmap='r')  
    plt.show()
S=test(200,1000)
print(pro_test(S))
print_test(S)
draw_test(S)   
x=[i*100 for i in range(1,21)]
y=[pro_test(test(500,j)) for j in x]
plt.plot(x, y, 'r.')
plt.ylim(0,1)
plt.xlabel("number of iterations")
plt.ylabel("proportion")
plt.show()